import backtrader as bt
import pandas as pd
import datetime


class DT_line(bt.Indicator):
    lines = ('U', 'D')
    params = (('period', 2), ('k_u', 0.7), ('k_d', 0.7))

    def __init__(self):
        self.addminperiod(self.p.period+1)

    def next(self):
        HH = max(self.data.high.get(-1, size=self.p.period))
        LC = min(self.data.close.get(-1, size=self.p.period))
        HC = max(self.data.close.get(-1, size=self.p.period))
        LL = min(self.data.low.get(-1, size=self.p.period))
        R = max(HH-LC, HC-LL)
        self.lines.U[0] = self.data.open[0] + self.p.k_u * R
        self.lines.D[0] = self.data.open[0] - self.p.k_d * R


class DualThrust(bt.Strategy):

    params = (('period', 2), ('k_u', 0.7), ('k_d', 0.7))

    def __init__(self):
        self.dataclose = self.data0.close  # self.data0 等于 maindata
        self.D_line = DT_line(self.data1,
                              period=self.p.period,
                              k_u=self.p.k_u,
                              k_d=self.p.k_d)
        self.D_line = self.D_line()  # 把日线图映射到分钟图的范围上
        self.D_line.plotinfo.plotmaster = self.data0  # 把DT_line与日线图合在一块

        # 当价格上穿上轨线时，做多
        self.buy_signal = bt.indicators.CrossOver(self.dataclose, self.D_line.U)  # 设置买入信号
        # 当价格下船下轨线时，做空
        self.sell_signal = bt.indicators.CrossDown(self.dataclose, self.D_line.D)  # 设置卖出信号

    def start(self):
        print("the word call me")

    def prenext(self):
        print("not mature")

    def next(self):

        if self.data.datetime.time()>datetime.time(0,2) \
            and self.data.datetime.time()< datetime.time(23,58):

            # 判断是否有持仓
            if not self.position and self.buy_signal[0] == 1:
                self.order = self.buy()
            if not self.position and self.sell_signal[0] == 1:
                self.order = self.sell()
            if self.getposition().size<0 and self.buy_signal[0] == 1:
                self.order = self.close()
                self.order = self.buy()
            if self.getposition().size>0 and self.sell_signal[0] == 1:
                self.order = self.close()
                self.order = self.sell()
        if self.data.datetime.time() >= datetime.time(23,58) and self.position:
            self.order = self.close()

    def stop(self):
        print("period: %s, k_u: %s, k_d: %s, final_value: %.2f" %
              (self.p.period, self.p.k_u, self.p.k_d, self.broker.getvalue()))

if __name__ == '__main__':
    cerebro = bt.Cerebro()

    dataframe = pd.read_csv("data/ETHUSD/bitmex.csv")
    dataframe['open_time'] = pd.to_datetime(dataframe['open_time'])
    dataframe.set_index('open_time', inplace=True)
    brf_min_bar = bt.feeds.PandasData(
        dataname = dataframe,
        fromdate=datetime.datetime(2019,10,15),
        todate=datetime.datetime(2019,10,30),
        timeframe=bt.TimeFrame.Minutes
    )

    # add Data
    cerebro.adddata(brf_min_bar)
    cerebro.resampledata(brf_min_bar, timeframe=bt.TimeFrame.Days)

    # add strategy
    # cerebro.addstrategy(DualThrust)
    cerebro.optstrategy(
        DualThrust,
        period = range(1,5),
        k_u=[n/10.0 for n in range(2,10)],
        k_d=[n/10.0 for n in range(2,10)]
    )

    cerebro.run()
    # cerebro.plot()
